package com.spark.test

import kafka.common.TopicAndPartition
import kafka.message.MessageAndMetadata
import kafka.serializer.StringDecoder
import kafka.utils.{ZKGroupTopicDirs, ZkUtils}
import org.I0Itec.zkclient.ZkClient
import org.apache.spark.SparkConf
import org.apache.spark.streaming.dstream.{DStream, InputDStream}
import org.apache.spark.streaming.kafka.{HasOffsetRanges, KafkaUtils, OffsetRange}
import org.apache.spark.streaming.{Duration, StreamingContext}
/**
  * 直连方式 kafka
  */
object KafkaDirctWordCount {

  def main(args: Array[String]): Unit = {
    // 指定组名
    val group : String = "g001"
    // 创建SparkConf
    val conf : SparkConf= new SparkConf().setAppName("KafkaDirctWordCount").setMaster("local[2]")
    // 创建StreamingContext,并设置间隔时间
    val ssc: StreamingContext = new StreamingContext(conf,Duration(5000))
    // 指定消费者的topic名字
    val topic = "test"
    // 指定kafka的broker地址(sparkStream的Task直连到kafka的分区上)
    val brokerList = "localhost:9092"

    // 指定zk地址 后期更新消息的偏移量时使用(以后可以使用redis,mysql记录偏移量)
    val zkQuorum = "localhost:2181"
    // 创建stream时使用topic名字集合，SparkStreaming可同时消费多个topic
    val topics : Set[String] = Set(topic)

    // 创建一个 ZKGroupTopicDirs对象，其实是指定往zk写入的目录，用于保存偏移量
    val topicDirs = new ZKGroupTopicDirs(group,topic)
    // 获取 zookeeper中的路径 "/g001/offsets/wordcount/"
    val zkTopicPath = s"${topicDirs.consumerOffsetDir}"

    // 准备kafka的参数
    val kafkaParams = Map(
      "metadata.broker.list" -> brokerList,
      "group.id" -> group,
      "auto.offset.reset" -> kafka.api.OffsetRequest.SmallestTimeString
    )

    // zookeeper的host和ip,创建一个client,用于更新偏移量
    // 是zookeeper的客户端,可以从zk中读取偏移量数据,并更新偏移量
    val zkClient = new ZkClient(zkQuorum)

    //查询该路径下是否字节点（默认有字节点为我们自己保存不同 partition 时生成的）
    // /g001/offsets/wordcount/0/10001"
    // /g001/offsets/wordcount/1/30001"
    // /g001/offsets/wordcount/2/10001"
    //zkTopicPath  -> /g001/offsets/wordcount/
    println("zkTopicPath "+zkTopicPath)
    val children = zkClient.countChildren(zkTopicPath)

    var kafkaStream : InputDStream[(String,String)] = null

    // 如果zk中有保存offset，我们会利用这个offset作为kafkaStream的起始位置
    var fromOffsets : Map[TopicAndPartition,Long] = Map()

    // 如果保存 offset
    if(children > 0){
//      // 保存过
      for(i <- 0 until children){
        // /g001/offsets/wordcount/0/10001
        // /g001/offsets/wordcount/0
        val partitionOffset = zkClient.readData[String](s"$zkTopicPath/${i}")
        println("partitionOffset "+partitionOffset)
        // wordcount/0
        val tp = TopicAndPartition(topic,i)
        //将不同 partition 对应的 offset 增加到 fromOffsets 中
        // wordcount/0 -> 10001
        fromOffsets += (tp -> partitionOffset.toLong)
      }
      //Key: kafka的key   values: "hello tom hello jerry"
      //这个会将 kafka 的消息进行 transform，最终 kafak 的数据都会变成 (kafka的key, message) 这样的 tuple
      val messageHandler = (mmd : MessageAndMetadata[String,String]) => (mmd.key(),mmd.message())
      //通过KafkaUtils创建直连的DStream（fromOffsets参数的作用是:按照前面计算好了的偏移量继续消费数据）
      //[String, String, StringDecoder, StringDecoder,     (String, String)]
      //  key    value    key的解码方式   value的解码方式
      kafkaStream = KafkaUtils.createDirectStream[String, String, StringDecoder, StringDecoder, (String, String)](ssc, kafkaParams, fromOffsets, messageHandler)
    } else{
      // 未保存过,根据kafkaParam的配置使用最新(largest)或者旧的(smallest) offset
      kafkaStream = KafkaUtils.createDirectStream[String,String,StringDecoder,StringDecoder](ssc,kafkaParams,topics);
    }

    println(kafkaStream)

    // 偏移量的范围
    var offsetRanges = Array[OffsetRange]()

    // 从kafka读取消息,DStream的Transform方法可以将当前批次的RDD获取出来
    // 该transform方法计算获取当前批次的RDD,然后将RDD的偏移量取出来,然后在将RDD返回到DStream
    val transform : DStream[(String,String)] = kafkaStream.transform{ rdd =>
      //得到该 rdd 对应 kafka 的消息的 offset
      //该RDD是一个KafkaRDD，可以获得偏移量的范围
      offsetRanges = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
      rdd
    }

    val messages: DStream[String] = transform.map(_._2)

    // 依次迭代DStream中的RDD
    messages.foreachRDD{ rdd =>
      println("rdd : "+rdd)
      rdd.foreachPartition { partition =>
        println("partition : "+partition)
        partition.foreach(x => {
          println(x)
        })
      }

      for ( o <- offsetRanges){
        //  /g001/offsets/wordcount/0
        val zkPath = s"${topicDirs.consumerOffsetDir}/${o.partition}"
        //将该 partition 的 offset 保存到 zookeeper
        //  /g001/offsets/wordcount/0/20000
        println(zkPath)
        println(o.untilOffset.toString)
        ZkUtils.updatePersistentPath(zkClient,zkPath,o.untilOffset.toString)
      }

    }


    ssc.start()
    ssc.awaitTermination()

  }
}
